#include "three_nn_n_d.h"

using namespace ThreeNN;

extern "C" __global__ __aicore__ void three_nn(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices,
                                               GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tilingData, tiling);
    if (TILING_KEY_IS(101)) {
        ThreeNNND<DTYPE_XYZ1, 1> op;
        op.Init(xyz1, xyz2, dist, indices, &tilingData);
        op.ProcessSmallX2();
    } else if (TILING_KEY_IS(201)) {
        ThreeNNND<DTYPE_XYZ1, 2> op;
        op.Init(xyz1, xyz2, dist, indices, &tilingData);
        op.ProcessBigX2();
    }
}